import time
from feature_extractor import Gauss2DFeatureExtractor
from mpl_toolkits.mplot3d import Axes3D  # 必须导入的三维工具包
from problem import Poisson_2d, Wave_2d, Diffusion_2d, Poisson_complex_2d
import numpy as np
from utils import  Domain, plot_tri_contourf, sample_boundaries
import matplotlib.pyplot as plt
from regression import LinearRegression
from options import Options

class Solver:
    def __init__(self, args):
        self.domain_range = args.domain
        self.problem_name = args.problem_name
        self.case = args.case
        self.k = args.k
        self.n_gaussian = args.n_gaussian
        self.n_level = args.n_level
        self.n_pde = args.n_pde
        self.n_bc = args.n_bc
        self.sigma = args.sigma
        self.mode = args.mode

        # seed
        np.random.seed(1234)

        # Set problem, extractor, regressor
        self.problem = self._problem()
        self.extractor = Gauss2DFeatureExtractor(n_gaussian=self.n_gaussian, n_level=self.n_level, sigma=self.sigma, mode=self.mode, domain=self.domain_range)
        print("Loading extractor2D")
        print(f"{self.extractor}")
        self.regressor = LinearRegression()


        # domain and sample
        self.domain = Domain(domain_range=self.domain_range, domain_shape='box')
        self.x_pde = self.domain.sample_domain_uniform(self.n_pde)
        # self.x_bc = self.domain.sample_boundary_uniform(self.n_bc)
        # 为了适用于更多问题，将四边的边界条件都表示出来
        self.x_top, self.x_bottom, self.x_left, self.x_right = sample_boundaries(domain_range=self.domain_range, num_samples=self.n_bc)
        self.x_bc = np.vstack([self.x_top, self.x_bottom, self.x_left, self.x_right])

    def _problem(self):
        assert self.problem_name in ["poisson2d", "wave2d", "diffusion2d", "poisson_complex2d"], f"need select from [poisson2d, wave2d, diffusion2d, poisson_complex2d]"
        if self.problem_name == "poisson2d":
            return Poisson_2d(case=self.case, k=self.k, domain=self.domain_range)
        elif self.problem_name == "wave2d":
            return Wave_2d(case=self.case, k=self.k, domain=self.domain_range)
        elif self.problem_name == "diffusion2d":
            return Diffusion_2d(case=self.case, k=self.k, domain=self.domain_range)
        elif self.problem_name == "poisson_complex2d":
            return Poisson_complex_2d(case=self.case, k=self.k, domain=self.domain_range)

    def get_features_target(self):
        if self.problem_name == "poisson2d":
            return self.get_features_target_poisson2d()
        elif self.problem_name == "wave2d":
            return self.get_features_target_wave2d()
        elif self.problem_name == "diffusion2d":
            return self.get_features_target_diffusion2d()
        elif self.problem_name == "poisson_complex2d":
            return self.get_features_target_poisson_complex2d()

    def get_features_target_poisson2d(self):
        '''
        u_xx + u_yy = f
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u', 'u00', 'u11'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u'))

        features_pde = z_pde['u00'] + z_pde['u11']
        features_bc = z_bc['u']
        
        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.bc(self.x_bc)
        target = np.concatenate([target_pde, target_bc])

        return features, target
    
    def get_features_target_wave2d(self):
        '''
            参考PIRBN  没有右边界条件，左边界条件有两个
            u_xx - 4u_yy = 0   默认[0, 1] x [0, 1] 
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u', 'u00', 'u11'))
        z_bc_top = self.extractor.eval(self.x_top, eval_list=('u'))
        z_bc_bottom = self.extractor.eval(self.x_bottom, eval_list=('u'))
        z_bc_left = self.extractor.eval(self.x_left, eval_list=('u', 'u0'))

        features_pde = z_pde['u00'] - 4 * z_pde['u11']
        features_bc_top = z_bc_top['u']
        features_bc_bottom = z_bc_bottom['u']
        features_bc_left_1, features_bc_left_2  = z_bc_left['u'],  z_bc_left['u0']
        
        features = np.vstack([features_pde, features_bc_top, features_bc_bottom, features_bc_left_1, features_bc_left_2])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc_top = self.problem.u_exact(self.x_top)
        target_bc_bottom = self.problem.u_exact(self.x_bottom)
        target_bc_left_1, target_bc_left_2 = self.problem.u_exact(self.x_left), self.problem.u_x_exact(self.x_left)
        target = np.concatenate([target_pde, target_bc_top, target_bc_bottom, target_bc_left_1, target_bc_left_2])

        return features, target
    
    def get_features_target_diffusion2d(self):
        '''
            参考PIRBN  没有上边界条件，其他全是狄利克雷边界条件
            u_t - 0.01u_xx = f(x)   默认[5, 10] x [5, 10] 
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u', 'u00', 'u1'))
        z_bc_bottom = self.extractor.eval(self.x_bottom, eval_list=('u'))
        z_bc_left = self.extractor.eval(self.x_left, eval_list=('u'))
        z_bc_right = self.extractor.eval(self.x_right, eval_list=('u'))

        features_pde = z_pde['u1'] - 0.01 * z_pde['u00']
        features_bc_bottom = z_bc_bottom['u']
        features_bc_left  = z_bc_left['u']
        features_bc_right = z_bc_right['u']
        
        features = np.vstack([features_pde, features_bc_bottom, features_bc_left, features_bc_right])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc_bottom = self.problem.u_exact(self.x_bottom)
        target_bc_left = self.problem.u_exact(self.x_left)
        target_bc_right = self.problem.u_exact(self.x_right)
        target = np.concatenate([target_pde, target_bc_bottom, target_bc_left, target_bc_right])

        return features, target
    
    def get_features_target_poisson_complex2d(self):
        '''
            参考RRNN的例子  多尺度椭圆方程
            -div(a_eps \nabla u) = f
        '''
        z_pde = self.extractor.eval(self.x_pde, eval_list=('u', 'u0', 'u1', 'u00', 'u11'))
        z_bc = self.extractor.eval(self.x_bc, eval_list=('u'))
        a_eps = self.problem.a_eps(self.x_pde)
        a_eps_x = self.problem.a_eps_x(self.x_pde)
        a_eps_y = self.problem.a_eps_y(self.x_pde)


        features_pde = - a_eps_x * z_pde['u0'] - a_eps * z_pde['u00'] - a_eps_y * z_pde['u1'] - a_eps * z_pde['u11'] 
        # 注意case3是poisson-Boltzmann方程
        if self.case == 3:
            features_pde = features_pde + np.pi * np.pi * z_pde['u']
        features_bc = z_bc['u']
        
        features = np.vstack([features_pde, features_bc])

        target_pde = self.problem.rhs(self.x_pde)
        target_bc = self.problem.u_exact(self.x_bc)
        target = np.concatenate([target_pde, target_bc])

        return features, target
    
    def plot_exact_solution_3d(self):
        X = np.linspace(0, 1, 100)
        Y = np.linspace(0, 1, 100)
        X, Y = np.meshgrid(X, Y)
        t = np.column_stack((X.flatten(), Y.flatten()))
        z = self.problem.a_eps(t)
        Z = z.reshape((100, 100))

        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection='3d')  # 设置投影类型为三维
        # 绘图
        surf = ax.plot_surface(X, Y, Z, cmap='jet')  # 透明度设置

        # 设置标签和标题
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('A')
        # ax.set_title('3D Surface Plot')

        # 添加颜色条
        fig.colorbar(surf)
        plt.show()

    def show_sample_points(self):
        plt.scatter(self.x_pde[:, 0], self.x_pde[:, 1], color='blue', marker='o', s=20)
        plt.show()
        

    def solve(self, plot=True):
        self.plot_exact_solution_3d()
        tt = time.time()
        # 1. features and target
        features, target = self.get_features_target()

        # 2. Least square
        self.regressor.fit(features, target)

        # 3. Test
        # plt.plot(self.regressor.w)
        x = self.domain.sample_domain_uniform(mesh_size=[100, 100])
        u = self.problem.u_exact(x)

        z = self.extractor.eval(x, eval_list=('u'))

        features = z['u']

        u_pred = self.regressor.predict(features)
        rmse = np.sqrt(np.mean((u - u_pred) ** 2))
        mae = np.max(np.abs(u - u_pred))
        print(f"problem_name: {self.problem_name}, case: {self.case}, Gaussians: {self.n_gaussian:4d}, "
              f"sigma: {self.sigma:.3f}, MSE: {rmse:.2e}, MAE: {mae:.2e}, "
              f"SumGaussians:{self.extractor.getSumGaussians()}, time:{time.time()-tt: .2f}")
        
        error = np.abs(u - u_pred)
        plot_tri_contourf(x[:, 0], x[:, 1], u[:, 0], title='exact solution')
        plot_tri_contourf(x[:, 0], x[:, 1], u_pred[:, 0], title='predicted solution')
        plot_tri_contourf(x[:, 0], x[:, 1], error[:, 0], title='error')
        plt.show()

        return rmse, mae


if __name__ == "__main__":
    # 2d problem
    args = Options().parse()
    args.domain = (0, 1, 0, 1) 
    args.problem_name = 'poisson_complex2d'  # ['poisson2d', 'wave2d', 'diffusion2d', 'poisson_complex2d']
    args.case = 1
    args.k = 0.2           # some problem case will use  np.sqrt(9200)  0.005  0.05
    args.n_gaussian = 300      # sigma_1
    args.n_level = 5
    args.n_pde = [50, 50]     # list
    args.n_bc = 200
    args.sigma = 0.5
    print("#"*10 + "args information" + "#"*10)
    print(args)

    solver = Solver(args)
    solver.solve()

